#!/usr/bin/python
# This script downloads complete genomes (some with multiple replicons)
# from NCBI website. Takes a tab-delimited file in the following format:
# 57995	Prochlorococcus marinus subsp. marinus str. CCMP1375	NC_005042
# 57761	Prochlorococcus marinus subsp. pastoris str. CCMP1986	NC_005072
# 57767	Nostoc punctiforme PCC 73102	NC_010628,NC_010632,NC_010630,NC_010633,NC_010629,NC_010631
# 57773	Prochlorococcus marinus str. MIT 9313	NC_005071
# 61581	Synechococcus sp. WH 8102	NC_005070
# 1st column is Refseq ID, 2nd column is organism name, and 3rd column is refseq accession numbers
# Usage: dissertation_DownloadGenomes.py todownload_cyano.txt

import sys
import re
from Bio import SeqIO
from Bio import Entrez
from Bio.Seq import Seq
from Bio.SeqRecord import SeqRecord
from Bio.Alphabet import generic_protein

listfile = sys.argv[1]
genomelist = open(listfile, "rU")
gl = genomelist.readlines()

Entrez.email = "jimmysaw@gmail.com"

for line in gl:
    l = line.split('\t')
    refseqid = l[0]
    organism = l[1]
    rsacc = l[2].rstrip()
    tmpline = rsacc.split(',')
    seqrecs = []
    seqfaas = []
    print "Working on" + "\t" + organism + "\t" + refseqid
    if len(tmpline) > 1:
        for t in tmpline:
            rec = Entrez.efetch(db="sequences", id=t, retmode="gb", rettype="text")
            seqs = SeqIO.read(rec, "gb")
            seqrecs.append(seqs)
            for feat in seqs.features:
                if feat.type == "CDS":
                    if feat.qualifiers.has_key('locus_tag'):
                        aa = feat.qualifiers['translation'][0]
                        aaid = feat.qualifiers['locus_tag'][0]
                        aadesc = feat.qualifiers['product'][0]
                        tmprec = SeqRecord(Seq(aa, generic_protein), id=aaid, description=aadesc)
                        seqfaas.append(tmprec)
    else:
        t = tmpline[0]
        rec = Entrez.efetch(db="sequences", id=t, retmode="gb", rettype="text")
        seq = SeqIO.read(rec, "gb")
        seqrecs.append(seq)
        for feat in seq.features:
            if feat.type == "CDS":
                if feat.qualifiers.has_key('locus_tag'):
                    aa = feat.qualifiers['translation'][0]
                    aaid = feat.qualifiers['locus_tag'][0]
                    aadesc = feat.qualifiers['product'][0]
                    tmprec = SeqRecord(Seq(aa, generic_protein), id=aaid, description=aadesc)
                    seqfaas.append(tmprec)
    sequences = SeqIO.parse(rec, 'fasta')
    outgbk = refseqid + ".refseq.gbk"
    outfaa = refseqid + ".refseq.faa"
    SeqIO.write(seqrecs, outgbk, "gb")
    SeqIO.write(seqfaas, outfaa, "fasta")

    print "Total number of CDS = ", len(seqfaas)
    print "Done with" + "\t" + organism + "\t" + refseqid

genomelist.close()

